"""
Database executor for executing SQL statements against KWDB.
"""
from __future__ import annotations

import logging
import io
import re
from typing import Optional, List, Tuple, Dict
from urllib.parse import urlparse, parse_qs
import psycopg2
from psycopg2.extras import execute_batch, execute_values
from psycopg2.pool import ThreadedConnectionPool

logger = logging.getLogger(__name__)


def parse_connection_string(conn_string: str) -> dict:
    """
    Parse PostgreSQL connection string.
    
    Format: postgresql://[user[:password]@][host][:port][/database][?params]
    Example: postgresql://root@127.0.0.1:26257/demo?sslmode=disable
    
    Returns:
        Dictionary with connection parameters
    """
    parsed = urlparse(conn_string)
    
    # Handle password: if parsed.password is None, it means no password in URL
    # But parsed.password can also be empty string if user@host format
    password = ''
    if '@' in parsed.netloc:
        # Check if password is explicitly provided (user:password@host)
        auth_part = parsed.netloc.rsplit('@', 1)[0]
        if ':' in auth_part:
            password = auth_part.split(':', 1)[1]
    
    result = {
        'host': parsed.hostname or 'localhost',
        'port': parsed.port or 26257,
        'user': parsed.username or 'root',
        'password': password,
        'database': parsed.path.lstrip('/') or 'defaultdb',
    }
    
    # Parse query parameters
    if parsed.query:
        params = parse_qs(parsed.query)
        # Handle sslmode
        if 'sslmode' in params:
            result['sslmode'] = params['sslmode'][0]
        else:
            result['sslmode'] = 'disable'
    else:
        result['sslmode'] = 'disable'
    
    return result


class KWDBExecutor:
    """Executor for KWDB database operations."""
    
    def __init__(
        self,
        connection_string: Optional[str] = None,
        host: Optional[str] = None,
        port: Optional[int] = None,
        user: Optional[str] = None,
        password: Optional[str] = None,
        database: Optional[str] = None,
        min_connections: int = 2,
        max_connections: int = 10,
    ):
        """
        Initialize KWDB executor with connection pool.
        
        Args:
            connection_string: PostgreSQL connection string (takes precedence over individual params)
                Format: postgresql://user:password@host:port/database?sslmode=disable
            host: Database host (ignored if connection_string provided)
            port: Database port (ignored if connection_string provided)
            user: Database user (ignored if connection_string provided)
            password: Database password (ignored if connection_string provided)
            database: Database name (ignored if connection_string provided)
            min_connections: Minimum pool size
            max_connections: Maximum pool size
        """
        # Parse connection string if provided
        if connection_string:
            params = parse_connection_string(connection_string)
            host = params['host']
            port = params['port']
            user = params['user']
            password = params['password']
            database = params['database']
            sslmode = params.get('sslmode', 'disable')
            self.connection_string = connection_string
        else:
            # Use individual parameters with defaults
            host = host or "localhost"
            port = port or 26257
            user = user or "root"
            password = password or ""
            database = database or "defaultdb"
            sslmode = 'disable'
            self.connection_string = (
                f"postgresql://{user}:{password}@{host}:{port}/{database}?sslmode=disable"
            )
        
        # Normalize and guard parameters
        try:
            port = int(port) if port is not None else 26257
        except Exception:
            port = 26257
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        # Auto/default pool sizes if None; ensure sane bounds
        if min_connections is None:
            min_connections = 2
        if max_connections is None:
            max_connections = max(10, int(min_connections) + 1)
        try:
            min_connections = int(min_connections)
            max_connections = int(max_connections)
        except Exception:
            min_connections = 2
            max_connections = 10
        if max_connections < min_connections:
            max_connections = max(min_connections, 2)
        self.min_connections = min_connections
        self.max_connections = max_connections
        
        # Connection pool with connection validation and keepalive
        try:
            self.pool = ThreadedConnectionPool(
                min_connections,
                max_connections,
                host=host,
                port=port,
                user=user,
                password=password,
                database=database,
                sslmode=sslmode,
            )
            logger.info(
                f"Created connection pool for KWDB at {host}:{port}/{database} "
                f"(min={min_connections}, max={max_connections})"
            )
        except Exception as e:
            logger.error(f"Failed to create connection pool: {e}")
            self.pool = None
            raise
    
    def _validate_connection(self, conn) -> bool:
        """Validate that a connection is still alive."""
        try:
            with conn.cursor() as cursor:
                cursor.execute("SELECT 1")
                cursor.fetchone()
            return True
        except Exception:
            return False
    
    def _get_valid_connection(self, retries: int = 3):
        """Get a valid connection from pool, retrying if connection is closed."""
        if not self.pool:
            raise RuntimeError("Connection pool not initialized")
        
        for attempt in range(retries):
            conn = self.pool.getconn()
            if self._validate_connection(conn):
                return conn
            else:
                # Connection is invalid, close it and try again
                try:
                    conn.close()
                except Exception:
                    pass
                # Put back invalid connection (pool will create new one)
                try:
                    self.pool.putconn(conn)
                except Exception:
                    pass
                if attempt < retries - 1:
                    logger.warning(f"Got invalid connection, retrying (attempt {attempt + 1}/{retries})")
        
        raise RuntimeError("Failed to get valid connection after retries")
    
    def execute_sql(self, sql: str) -> None:
        """
        Execute a single SQL statement.
        
        Args:
            sql: SQL statement to execute
        """
        # Retry on transient errors like server restart/connection closed
        max_attempts = 5
        backoff_seconds = 1.0
        last_error = None
        for attempt in range(1, max_attempts + 1):
            conn = None
            try:
                conn = self._get_valid_connection()
                with conn.cursor() as cursor:
                    cursor.execute(sql)
                conn.commit()
                # Success
                if conn:
                    try:
                        self.pool.putconn(conn)
                    except Exception:
                        try:
                            conn.close()
                        except Exception:
                            pass
                return
            except Exception as e:
                last_error = e
                # Rollback if possible
                if conn:
                    try:
                        conn.rollback()
                    except Exception:
                        pass
                    try:
                        self.pool.putconn(conn)
                    except Exception:
                        try:
                            conn.close()
                        except Exception:
                            pass
                # Determine if retryable
                err_text = str(e).lower()
                retryable = any(
                    key in err_text for key in [
                        "server closed the connection",
                        "connection not open",
                        "connection reset",
                        "connection refused",
                        "could not send data to server",
                        "terminating connection",
                    ]
                )
                if attempt < max_attempts and retryable:
                    logger.warning(
                        f"Transient DB error on attempt {attempt}/{max_attempts}: {e}. Retrying after {backoff_seconds:.1f}s..."
                    )
                    import time
                    time.sleep(backoff_seconds)
                    # Exponential backoff up to a cap
                    backoff_seconds = min(backoff_seconds * 2.0, 8.0)
                    continue
                # Non-retryable or max attempts exceeded
                try:
                    from .diagnostics import diagnose_sql_error
                    diag = diagnose_sql_error(sql, str(e), phase='DML' if sql.strip().upper().startswith('INSERT') else 'DDL')
                except Exception:
                    diag = f"Error executing SQL: {e}\nSQL: {sql[:200]}"
                # Collect diagnostic error instead of logging immediately
                if not hasattr(self, '_diagnostic_errors'):
                    self._diagnostic_errors = []
                self._diagnostic_errors.append(diag)
                # Still log a brief error message
                logger.error(f"SQL execution failed: {e}")
                raise RuntimeError(diag)
    
    def execute_sql_batch(self, sql_statements: list[str]) -> None:
        """
        Execute multiple SQL statements in a single transaction.
        Optimized for batch execution with proper error handling.
        
        Args:
            sql_statements: List of SQL statements to execute
        """
        if not sql_statements:
            return
        
        conn = None
        try:
            conn = self._get_valid_connection()
            # Use a single transaction for all statements
            with conn.cursor() as cursor:
                # Execute all statements in one transaction
                for sql in sql_statements:
                    cursor.execute(sql)
            
            conn.commit()
        except Exception as e:
            if conn:
                try:
                    conn.rollback()
                except Exception:
                    # Connection may already be closed, ignore rollback error
                    pass
            try:
                from .diagnostics import diagnose_sql_error
                failing = ''
                # best-effort include last non-empty statement for context
                for s in reversed(sql_statements or []):
                    if s and s.strip():
                        failing = s
                        break
                diag = diagnose_sql_error(failing, str(e), phase='BATCH')
            except Exception:
                diag = f"Error executing SQL batch: {e}"
            # Collect diagnostic error instead of logging immediately
            if not hasattr(self, '_diagnostic_errors'):
                self._diagnostic_errors = []
            self._diagnostic_errors.append(diag)
            # Still log a brief error message
            logger.error(f"Batch execution failed: {e}")
            raise RuntimeError(diag)
        finally:
            if conn:
                try:
                    self.pool.putconn(conn)
                except Exception as e:
                    logger.warning(f"Error returning connection to pool: {e}")
                    # Connection may be invalid, try to close it
                    try:
                        conn.close()
                    except Exception:
                        pass
    
    def execute_ddl(self, ddl_sql: str) -> None:
        """
        Execute DDL statements (CREATE TABLE, etc.).
        
        Args:
            ddl_sql: DDL SQL statement(s)
        """
        statements = [s.strip() for s in ddl_sql.split(';') if s.strip()]
        for stmt in statements:
            if stmt:
                try:
                    self.execute_sql(stmt + ';')
                except Exception as e:
                    # e 已经是带诊断信息的 RuntimeError
                    raise
    
    def execute_constraints(self, constraints_sql: str) -> None:
        """
        Execute constraint/index creation statements.
        Optimized for memory usage - executes one at a time with progress logging.
        
        Args:
            constraints_sql: Constraint SQL statement(s)
        """
        statements = [s.strip() for s in constraints_sql.split(';') if s.strip()]
        total_statements = len(statements)
        
        if total_statements == 0:
            return
        
        logger.info(f"Executing {total_statements} constraint/index statements...")
        
        for idx, stmt in enumerate(statements, 1):
            if stmt:
                try:
                    # Log progress every 5 statements or for important operations
                    if idx % 5 == 0 or idx == 1 or idx == total_statements:
                        logger.info(f"Executing constraint/index {idx}/{total_statements}...")
                    
                    self.execute_sql(stmt + ';')
                    
                    # Small delay after each statement to allow memory cleanup
                    # Heavier delay for PK/INDEX which are more memory intensive (Cockroach/PG backfill)
                    if idx < total_statements:  # Don't delay after last statement
                        import time
                        heavy = bool(re.search(r"ADD\s+PRIMARY\s+KEY|CREATE\s+UNIQUE?\s*INDEX|CREATE\s+INDEX", stmt, re.IGNORECASE))
                        time.sleep(1.0 if heavy else 0.1)
                except Exception as e:
                    try:
                        from .diagnostics import diagnose_sql_error
                        diag = diagnose_sql_error(stmt, str(e), phase='CONSTRAINTS')
                    except Exception:
                        diag = f"Error executing constraint/index statement {idx}/{total_statements}: {e}\nFailed statement: {stmt[:200]}..."
                    # Collect diagnostic error instead of logging immediately
                    if not hasattr(self, '_diagnostic_errors'):
                        self._diagnostic_errors = []
                    self._diagnostic_errors.append(diag)
                    # Still log a brief error message
                    logger.error(f"Constraint execution failed: {e}")
                    raise RuntimeError(diag)
    
    def close(self):
        """Close connection pool."""
        if self.pool:
            try:
                self.pool.closeall()
                logger.info("Closed connection pool")
            except Exception as e:
                logger.warning(f"Error closing connection pool: {e}")
            finally:
                self.pool = None

    def fetch_scalar(self, sql: str):
        """Execute a SELECT and return the first column of the first row, or None."""
        conn = None
        try:
            conn = self._get_valid_connection()
            with conn.cursor() as cursor:
                cursor.execute(sql)
                row = cursor.fetchone()
                return row[0] if row else None
        finally:
            if conn:
                try:
                    self.pool.putconn(conn)
                except Exception:
                    try:
                        conn.close()
                    except Exception:
                        pass

    def fetch_table_row_counts(self, table_names: list[str]) -> dict[str, int]:
        """Return COUNT(*) for each table (quoted identifiers)."""
        results: dict[str, int] = {}
        conn = None
        try:
            conn = self._get_valid_connection()
            with conn.cursor() as cursor:
                for t in table_names:
                    try:
                        # 1) Try exact quoted identifier (case-sensitive)
                        cursor.execute(f'SELECT COUNT(*) FROM "{t}"')
                        row = cursor.fetchone()
                        results[t] = int(row[0]) if row and row[0] is not None else 0
                    except Exception:
                        # 2) Fallback: try unquoted (case-insensitive, folded to lower)
                        try:
                            cursor.execute(f'SELECT COUNT(*) FROM {t}')
                            row = cursor.fetchone()
                            results[t] = int(row[0]) if row and row[0] is not None else 0
                        except Exception:
                            # 3) Fallback: lowercased quoted
                            try:
                                tl = str(t).lower()
                                cursor.execute(f'SELECT COUNT(*) FROM "{tl}"')
                                row = cursor.fetchone()
                                results[t] = int(row[0]) if row and row[0] is not None else 0
                            except Exception:
                                results[t] = -1
        finally:
            if conn:
                try:
                    self.pool.putconn(conn)
                except Exception:
                    try:
                        conn.close()
                    except Exception:
                        pass
        return results

    def list_tables(self) -> list[str]:
        """List user tables in current database (Postgres-compatible catalogs)."""
        tables: list[str] = []
        conn = None
        try:
            conn = self._get_valid_connection()
            with conn.cursor() as cursor:
                cursor.execute(
                    """
                    SELECT table_name
                    FROM information_schema.tables
                    WHERE table_schema = 'public' AND table_type = 'BASE TABLE'
                    ORDER BY table_name
                    """
                )
                for r in cursor.fetchall() or []:
                    tables.append(r[0])
        finally:
            if conn:
                try:
                    self.pool.putconn(conn)
                except Exception:
                    try:
                        conn.close()
                    except Exception:
                        pass
        return tables

